#!/usr/bin/env python3
# _*_ coding:utf-8 _*_
'''
4. 使用光流估计方法，在前述测试视频上计算特征点，进一步进行特征点光流估计。(1)
'''
import cv2 as cv
# 加载视频
# Videofilename = r'D:\OpenCV-4.2.0-vc14_vc15\opencv\sources\samples\data\vtest.avi'
cap = cv.VideoCapture('vtest.avi')
if not cap.isOpened():
    print("无法打开视频文件")

# 角点检测参数设定
feature_params = dict(maxCorners=100,
                      qualityLevel=0.3,
                      minDistance=7,
                      blockSize=7)

# lucas_kanade光流法参数设定
lk_params = dict(winSize=(15, 15),
                 maxLevel=2,
                 criteria=(cv.TERM_CRITERIA_EPS | cv.TERM_CRITERIA_COUNT, 10, 0.03))

# 计算第一帧特征点
ret, prev = cap.read()                             # 读取视频帧
frame_gray = cv.cvtColor(prev, cv.COLOR_BGR2GRAY)  # 转化为灰度虚图像
p0 = cv.goodFeaturesToTrack(frame_gray, mask=None, **feature_params)  # 角点检测

while True:
    flag, source = cap.read()
    if not flag:                           # 未读到当前帧，结束
        break
    gray = cv.cvtColor(source, cv.COLOR_BGR2GRAY)
    # 计算光流
    p1, st, err = cv.calcOpticalFlowPyrLK(frame_gray, gray, p0, None,**lk_params)  # 前一帧的角点和当前帧的图像作为输入来得到角点在当前帧的位置
    # 获取好的特征点
    GoodPoints = p1[st==1]
    GoodPrePoints = p0[st == 1]
    # 在图像上画出光流向量
    res = source.copy()
    for i,(cur,pre) in enumerate(zip(GoodPoints,GoodPrePoints)):
        x0,y0 = cur.ravel()
        x1,y1 = pre.ravel()
        cv.line(res,(x0,y0),(x1,y1),(0,255,0))
        cv.circle(res,(x0,y0),3,(0,0,255))
    # 更新上一帧
    frame_gray = gray.copy()
    p0 = GoodPoints.reshape(-1,1,2)
    cv.imshow('lk_track', res)

    key = cv.waitKey(100)
    if key == 27:
       break